'''
-*- coding: utf-8 -*-
@File  : 01-demo.py
@Author: Shanmh
@Time  : 2024/03/06 下午4:03
@Function： openvino yolov8-cpu 推理
1座10核：Intel(R) Core(TM) i9-10900K CPU @ 3.70GHz
测试：
 1000张约 84fps


'''
import os.path
import time
import onnx

from ultralytics import YOLO
from typing import Tuple
from ultralytics.utils import ops
import torch
import numpy as np
import ipywidgets as widgets
import openvino as ov
from typing import Tuple, Dict
import cv2
import numpy as np

import random
import yaml
import xml.etree.ElementTree as ET



class Detection:
    def __init__(self,modelpath,yolov5_local,img_size=None):
        self.imgsz=[640,640] if img_size is None else img_size #(height, width)
        self.yolov5_local =yolov5_local
        self.modelpath=self.conv_(modelpath)
        self.colors = [(255, 0, 0), (128, 0, 0), (0, 255, 0), (0, 0, 255), (0, 0, 255)]


    def conv_(self,modelpath,half=True):
        if modelpath.endswith(".onnx"):
            return modelpath
        elif modelpath.endswith(".pt"):
            path_, name_ = os.path.split(modelpath)
            onnx_path=modelpath[:-2]+ "onnx"
            yaml_path=modelpath[:-2]+ "yaml"
            if not os.path.exists(onnx_path):
                det_model = torch.hub.load(self.yolov5_local, "custom", path = modelpath,
               source = 'local',force_reload=True)
                det_model.cpu()
                # if half:
                #     det_model.half()

                data = torch.zeros(1, 3, self.imgsz[0],self.imgsz[1])  # .cuda()
                torch.onnx.export(det_model, data,onnx_path,
                                  opset_version=11,
                                  export_params=True,
                                  verbose=True,
                                  input_names=['input'],
                                  output_names=['output'])
                #保存类别名字
                with open(yaml_path, 'w') as file:
                    yaml.dump(det_model.names, file)
            return onnx_path

    def init_(self):
        ie = ov.Core()
        ie.set_property({'CACHE_DIR': f'{self.modelpath[:-5]}_openvino_cache'})

        # devices = ie.available_devices
        ie.set_property({'CACHE_DIR': './openvino_cache'})
        self.model = ie.compile_model(model=ie.read_model(self.modelpath), device_name="CPU",compress_to_fp16=True)
        with open(self.modelpath[:-4]+"yaml", 'r') as file:
            categoryNames = yaml.safe_load(file)
        self.categoryNames = categoryNames

        self.input_key = self.model.input(0)
        self.output_key = self.model.output(0)
        self.networkInputHeight, self.networkInputWidth = list(self.input_key.shape)[2:]

        # if device.value != "CPU":
        #     self.det_ov_model .reshape({0: [1, 3, self.imgsz[0], self.imgsz[1]]})
        # self.det_compiled_model =  self.core.compile_model(self.det_ov_model, device.value)

    def detect(self, img,conf=0.25,iou=0.7,draw=True):
        resized_image = cv2.resize(src=img, dsize=(self.networkInputWidth, self.networkInputHeight))
        inputImage = np.expand_dims(np.transpose(resized_image, (2, 0, 1)), 0)
        inputImage = inputImage.astype(np.float32) / 255.0
        result = self.model([inputImage])[self.output_key]
        # print("model: ", time.time()-start)
        result = self.yolov5Post(img, result[0],conf,iou)
        if draw:
            img=self.draw(img,result)
        return img,result

    def draw(self, img, results):
        for r in results:
            x1, y1, x2, y2 = int(r["xmin"]), int(r["ymin"]), int(r["xmax"]), int(r["ymax"])
            score, c_class, name = r["confidence"], r["class"], r["name"]
            color = self.colors[-1] if c_class >= len(self.colors) or c_class < 0 else self.colors[c_class]
            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
            cv2.putText(img, '{}, {:.3f}'.format(name, score),
                                (x1, y1 + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                                color, 2)
        return img



    def yolov5Post(self, inputImage, outputData,conf,iou):
        # print(outputData.shape, type(outputData))
        imageHeight, imageWidth, _ = inputImage.shape

        xFactor = imageWidth / self.networkInputWidth
        yFactor = imageHeight / self.networkInputHeight

        confidence = outputData[:,4]
        outputData = outputData[confidence>=0.25]
        if len(outputData) == 0:
            return []
        classesScores = outputData[:,5:]
        classIds = np.argmax(classesScores, axis=1)
        rows = np.array(range(classesScores.shape[0]))

        indexs = classesScores[rows, classIds] > .25
        outputData = outputData[indexs]
        if len(outputData) == 0:
            return []
        confidences = outputData[:,4]
        boxes = outputData[:,:4]
        classIds = classIds[indexs]

        boxes[:,0] = (boxes[:,0] - 0.5 * boxes[:,2]) * xFactor
        boxes[:,1] = (boxes[:,1] - 0.5 * boxes[:,3]) * yFactor
        boxes[:,2] = boxes[:,2] * xFactor
        boxes[:,3] = boxes[:,3] * yFactor
        boxes = boxes.astype(np.int64)
        indexes = cv2.dnn.NMSBoxes(boxes, confidences, conf, iou)

        resultConfidences = confidences[indexes]
        resultClassIds = classIds[indexes]
        resultBoxes = boxes[indexes]

        results = []
        for i in range(len(resultBoxes)):
            r = {}
            r["confidence"] = resultConfidences[i]
            r["class"] = resultClassIds[i]
            r["xmin"] = resultBoxes[i,0]
            r["ymin"] = resultBoxes[i,1]
            r["xmax"] = resultBoxes[i,0] + resultBoxes[i,2]
            r["ymax"] = resultBoxes[i,1] + resultBoxes[i,3]
            r["name"] = "unknowedCategory" if resultClassIds[i] >= len(self.categoryNames) else self.categoryNames[resultClassIds[i]]
            results.append(r)
        return results


if __name__ == '__main__':
    detecter=Detection("/home/hxzh/PycharmProjects/CpuTest/models/yolov5s.pt","./yolov5")
    detecter.init_()
    # img=cv2.imread("/home/hxzh/Dataset/WaterDrop/dataset0311/images/train/20240311105903_water_194404.jpg")
    img=cv2.imread("/home/hxzh/PycharmProjects/CpuTest/asset/images/person.png")
    img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)

    st_time=time.time()
    test_count=1000
    for i in range(test_count):
        show,result=detecter.detect(img.copy(),0.25,0.7)

        cv2.imshow("a",show)
        cv2.waitKey(1)
    print(f"calculate 100 images spend {time.time()-st_time} ; speed: {test_count/(time.time()-st_time)}/s")

